import numpy as np
import pandas as pd
from qtorch.strategy import Strategy

class MomentumStrategy(Strategy):
    """动量策略（20日收益率）"""
    def __init__(self, window=20, threshold=0.05):
        self.window = window
        self.threshold = threshold
        
    def generate_signals(self, prices):
        returns = prices.pct_change(periods=self.window)
        signals = np.where(returns > self.threshold, 1,
                          np.where(returns < -self.threshold, -1, 0))
        return signals.astype(int)